from typing import Any, Dict, List, Tuple, TypeVar, Generic, Type, Union, NewType

from pydantic import BaseModel
from tortoise import Model
from tortoise.expressions import Q
from tortoise.transactions import in_transaction

from core.logger import logger

Total = NewType("Total", int)
ModelType = TypeVar("ModelType", bound=Model)
CreateSchemaType = TypeVar("CreateSchemaType", bound=BaseModel)
UpdateSchemaType = TypeVar("UpdateSchemaType", bound=BaseModel)

class CurdBase(Generic[ModelType, CreateSchemaType, UpdateSchemaType]):
    def __init__(self, model: Type[ModelType]):
        self.model = model
        
    def get_model_fields(self):
        values = [val["name"] for val in self.model.describe()["data_fields"]]
        return values
        
    async def get(self, **kwargs) -> ModelType:
        try:
            return await self.model.get(**kwargs)
        except:
            return None
    
    async def list(self,  page: int, page_size: int, search: Dict = dict(), order: list = []) -> Tuple[Total, List[ModelType]]:
        base_Q = Q()
        for k, v in search.items():
            base_Q &= Q(k=v)
        query = self.model.filter(base_Q)
        return await query.count(), await query.offset((page - 1) * page_size).limit(page_size).order_by(*order)
    
    async def create(self, obj_in: Union[UpdateSchemaType, Dict[str, Any]]) -> ModelType:
        if isinstance(obj_in, Dict):
            obj_dict = obj_in
        else:
            obj_dict = obj_in.dict()
        obj_dict = {k:v for k, v in obj_dict.items() if k in self.get_model_fields()}
        obj = self.model(**obj_dict)
        await obj.save()
        return obj

    async def update(self, id: int, obj_in: Union[UpdateSchemaType, Dict[str, Any]]) -> ModelType:
        if isinstance(obj_in, dict):
            obj_dict = obj_in
        else:
            obj_dict = obj_in.dict(exclude_unset=True, exclude={"id", "uid"})
        obj_dict = {k:v for k, v in obj_dict.items() if k in self.get_model_fields()}
        async with in_transaction() as connection:
            obj = await self.model.filter(id=id).using_db(connection).select_for_update().first()
            if obj is None:
                raise ValueError("Object not found")

            current_version = obj.version
            obj_dict['version'] = current_version + 1
            obj.update_from_dict(obj_dict)
            update_count = await self.model.filter(id=id, version=current_version).using_db(connection).update(**obj_dict)

            if update_count == 0:
                err_info = "Update conflict, the record has been modified by another transaction."
                logger.warning(f"{id} - {err_info}")
                raise ValueError(err_info)

            await obj.save(using_db=connection)
            return obj

    async def remove(self, **kwargs) -> None:
        obj = await self.get(**kwargs)
        if obj:
            await obj.delete()
        else:
            logger.warning(f"Not match object")